{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9-final"
    },
    "colab": {
      "name": "Robotic Manipulation - Force Control.ipynb",
      "provenance": [],
      "collapsed_sections": []
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EgiF12Hf1Dhs"
      },
      "source": [
        "# Notebook setup\n",
        "\n",
        "The following cell will:\n",
        "- on Colab (only), install Drake to `/opt/drake`, install Drake's prerequisites via `apt`, and add pydrake to `sys.path`.  This will take approximately two minutes on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  If you navigate between notebooks using Colab's \"File->Open\" menu, then you can avoid provisioning a separate machine for each notebook.\n",
        "- launch a server for our 3D visualizer (MeshCat) that will be used for the remainder of this notebook.\n",
        "\n",
        "You will need to rerun this cell if you restart the kernel, but it should be fast because the machine will already have drake installed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eeMrMI0-1Dhu"
      },
      "source": [
        "import importlib\n",
        "import sys\n",
        "from urllib.request import urlretrieve\n",
        "\n",
        "# Install drake.\n",
        "if 'google.colab' in sys.modules and importlib.util.find_spec('manipulation') is None:\n",
        "    urlretrieve(f\"http://manipulation.csail.mit.edu/scripts/setup/setup_manipulation_colab.py\",\n",
        "                \"setup_manipulation_colab.py\")\n",
        "    from setup_manipulation_colab import setup_manipulation\n",
        "    setup_manipulation(manipulation_sha='master', drake_version='20201120', drake_build='nightly')\n",
        "\n",
        "# Determine if this notebook is currently running as a notebook or a unit test.\n",
        "from IPython import get_ipython\n",
        "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
        "\n",
        "# Setup rendering (with xvfb), if necessary:\n",
        "import os\n",
        "if 'google.colab' in sys.modules and os.getenv(\"DISPLAY\") is None:\n",
        "    from pyvirtualdisplay import Display\n",
        "    display = Display(visible=0, size=(1400, 900))\n",
        "    display.start()\n",
        "\n",
        "# Use pyngrok on colab.\n",
        "server_args = []\n",
        "if 'google.colab' in sys.modules:\n",
        "  server_args = ['--ngrok_http_tunnel']\n",
        "\n",
        "# Start a single meshcat server instance to use for the remainder of this notebook.\n",
        "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
        "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
        "\n",
        "# Let's do all of our imports here, too.\n",
        "import numpy as np\n",
        "import ipywidgets\n",
        "import pydot\n",
        "import pydrake.all\n",
        "import os\n",
        "from IPython.display import display, SVG\n",
        "\n",
        "import pydrake.all"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "source": [
        "# Set up a basic ManipulationStation diagram\n",
        "\n",
        "Completely in pydrake.  Feel free to modify it as you see fit.  You might also look at the [C++ version](https://github.com/RobotLocomotion/drake/blob/master/examples/manipulation_station/manipulation_station.cc#L193) if you want inspiration for tables / cupboards / bins, etc that you might add.  Here is [a link to the scenarios file](https://github.com/RussTedrake/manipulation/blob/master/manipulation/scenarios.py), in case you need to modify `AddIiwa` and friends."
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from manipulation.scenarios import AddIiwa, AddWsg, AddRgbdSensors\n",
        "from manipulation.utils import FindResource\n",
        "\n",
        "def MakeManipulationStation(time_step=0.002):\n",
        "    builder = pydrake.systems.framework.DiagramBuilder()\n",
        "\n",
        "    # Add (only) the iiwa, WSG, and cameras to the scene.\n",
        "    plant, scene_graph = pydrake.multibody.plant.AddMultibodyPlantSceneGraph(\n",
        "        builder, time_step=time_step)\n",
        "    iiwa = AddIiwa(plant)\n",
        "    wsg = AddWsg(plant, iiwa)\n",
        "    pydrake.multibody.parsing.Parser(plant).AddModelFromFile(\n",
        "        FindResource(\"models/camera_box.sdf\"), \"camera0\")\n",
        "    plant.Finalize()\n",
        "\n",
        "    num_iiwa_positions = plant.num_positions(iiwa)\n",
        "\n",
        "    # I need a PassThrough system so that I can export the input port.\n",
        "    iiwa_position = builder.AddSystem(pydrake.systems.primitives.PassThrough(num_iiwa_positions))\n",
        "    builder.ExportInput(iiwa_position.get_input_port(), \"iiwa_position\")\n",
        "    builder.ExportOutput(iiwa_position.get_output_port(), \"iiwa_position_command\")\n",
        "\n",
        "    # Export the iiwa \"state\" outputs.\n",
        "    demux = builder.AddSystem(pydrake.systems.primitives.Demultiplexer(\n",
        "        2 * num_iiwa_positions, num_iiwa_positions))\n",
        "    builder.Connect(plant.get_state_output_port(iiwa), demux.get_input_port())\n",
        "    builder.ExportOutput(demux.get_output_port(0), \"iiwa_position_measured\")\n",
        "    builder.ExportOutput(demux.get_output_port(1), \"iiwa_velocity_estimated\")\n",
        "    builder.ExportOutput(plant.get_state_output_port(iiwa), \"iiwa_state_estimated\")\n",
        "\n",
        "    # Make the plant for the iiwa controller to use.\n",
        "    controller_plant = pydrake.multibody.plant.MultibodyPlant(time_step=time_step)\n",
        "    controller_iiwa = AddIiwa(controller_plant)\n",
        "    AddWsg(controller_plant, controller_iiwa, welded=True)\n",
        "    controller_plant.Finalize()\n",
        "\n",
        "    # Add the iiwa controller\n",
        "    iiwa_controller = builder.AddSystem(\n",
        "        pydrake.systems.controllers.InverseDynamicsController(\n",
        "            controller_plant,\n",
        "            kp=[100]*num_iiwa_positions,\n",
        "            ki=[1]*num_iiwa_positions,\n",
        "            kd=[20]*num_iiwa_positions,\n",
        "            has_reference_acceleration=False))\n",
        "    iiwa_controller.set_name(\"iiwa_controller\")\n",
        "    builder.Connect(\n",
        "        plant.get_state_output_port(iiwa), iiwa_controller.get_input_port_estimated_state())\n",
        "\n",
        "    # Add in the feed-forward torque\n",
        "    adder = builder.AddSystem(pydrake.systems.primitives.Adder(2, num_iiwa_positions))\n",
        "    builder.Connect(iiwa_controller.get_output_port_control(),\n",
        "                    adder.get_input_port(0))\n",
        "    # Use a PassThrough to make the port optional (it will provide zero values if not connected).\n",
        "    torque_passthrough = builder.AddSystem(\n",
        "        pydrake.systems.primitives.PassThrough([0]*num_iiwa_positions))\n",
        "    builder.Connect(torque_passthrough.get_output_port(), adder.get_input_port(1))\n",
        "    builder.ExportInput(torque_passthrough.get_input_port(), \"iiwa_feedforward_torque\")\n",
        "    builder.Connect(adder.get_output_port(), plant.get_actuation_input_port(iiwa))\n",
        "\n",
        "    # Add discrete derivative to command velocities.\n",
        "    desired_state_from_position = builder.AddSystem(\n",
        "        pydrake.systems.primitives.StateInterpolatorWithDiscreteDerivative(\n",
        "            num_iiwa_positions, time_step, suppress_initial_transient=True))\n",
        "    desired_state_from_position.set_name(\"desired_state_from_position\")\n",
        "    builder.Connect(desired_state_from_position.get_output_port(),      \n",
        "                    iiwa_controller.get_input_port_desired_state())\n",
        "    builder.Connect(iiwa_position.get_output_port(), desired_state_from_position.get_input_port())\n",
        "\n",
        "    # Export commanded torques.\n",
        "    #builder.ExportOutput(adder.get_output_port(), \"iiwa_torque_commanded\")\n",
        "    #builder.ExportOutput(adder.get_output_port(), \"iiwa_torque_measured\")\n",
        "\n",
        "    # Wsg controller.\n",
        "    wsg_controller = builder.AddSystem(pydrake.manipulation.schunk_wsg.SchunkWsgPositionController())\n",
        "    wsg_controller.set_name(\"wsg_controller\")\n",
        "    builder.Connect(wsg_controller.get_generalized_force_output_port(),             \n",
        "                    plant.get_actuation_input_port(wsg))\n",
        "    builder.Connect(plant.get_state_output_port(wsg), wsg_controller.get_state_input_port())\n",
        "    builder.ExportInput(wsg_controller.get_desired_position_input_port(), \"wsg_position\")\n",
        "    builder.ExportInput(wsg_controller.get_force_limit_input_port(), \"wsg_force_limit\")\n",
        "    wsg_mbp_state_to_wsg_state = builder.AddSystem(\n",
        "        pydrake.manipulation.schunk_wsg.MakeMultibodyStateToWsgStateSystem())\n",
        "    builder.Connect(plant.get_state_output_port(wsg), wsg_mbp_state_to_wsg_state.get_input_port())\n",
        "    builder.ExportOutput(wsg_mbp_state_to_wsg_state.get_output_port(), \"wsg_state_measured\")\n",
        "    builder.ExportOutput(wsg_controller.get_grip_force_output_port(), \"wsg_force_measured\")\n",
        "\n",
        "    # Cameras.\n",
        "    AddRgbdSensors(builder, plant, scene_graph)\n",
        "\n",
        "    # Export \"cheat\" ports.\n",
        "    builder.ExportOutput(scene_graph.get_query_output_port(), \"geometry_query\")\n",
        "    builder.ExportOutput(plant.get_contact_results_output_port(), \"contact_results\")\n",
        "    builder.ExportOutput(plant.get_state_output_port(), \"plant_continuous_state\")\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    return diagram\n",
        "\n",
        "diagram = MakeManipulationStation()\n",
        "\n",
        "display(SVG(pydot.graph_from_dot_data(diagram.GetGraphvizString())[0].create_svg()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def TestWithMeshcat():\n",
        "    builder = pydrake.systems.framework.DiagramBuilder()\n",
        "    station = builder.AddSystem(MakeManipulationStation())\n",
        "\n",
        "    visualizer = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(\n",
        "        builder, output_port=station.GetOutputPort(\"geometry_query\"), zmq_url=zmq_url)\n",
        "\n",
        "    diagram = builder.Build()\n",
        "\n",
        "    context = diagram.CreateDefaultContext()\n",
        "    visualizer.load(visualizer.GetMyContextFromRoot(context))\n",
        "    diagram.Publish(context)\n",
        "\n",
        "TestWithMeshcat()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ]
}